package com.starhub.application.collection.common;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.starhub.application.collection.dto.DataDefinitionDto;
import com.starhub.application.collection.dto.MilvusCollectionParamDto;
import com.starhub.application.model.enums.VectorModelTypeEnum;
import com.starhub.application.model.util.DimensionUtil;
import com.starhub.application.rag.biz.MilvusClientFactory;
import com.starhub.application.rag.constant.RagConstant;

import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.index.CreateIndexParam;
import lombok.extern.slf4j.Slf4j;

/**
 * Milvus集合创建类
 * @version: V1.0
 * @author: liuhf
 * @time 2025-06-06
 */
@Slf4j
@Component
public class MilvusCreationHelper {
    
    @Autowired
    private MilvusClientFactory milvusClientFactory;

    /**
     * 创建默认的集合
     * 
     * @param collection
     * @return
     */
    public void createDefaultCollection(DataDefinitionDto collection) {
        MilvusServiceClient milvusClient = null;
        try {
            milvusClient = milvusClientFactory.getClient(RagConstant.MILVUS_MARK);

            MilvusCollectionParamDto milvusCollectionParamDto = getDefaultCreateCollectionParam(
                    collection.getCollectionIdentifier(), collection.getDescription(),
                    collection.getVectorModel());

            milvusClient.createCollection(milvusCollectionParamDto.getCreateCollectionParam());
            milvusClient.createIndex(milvusCollectionParamDto.getCreateIndexParam());
            
            log.info("成功创建Milvus集合: {}", collection.getCollectionIdentifier());
        } catch (Exception e) {
            log.error("创建Milvus集合失败: {}", collection.getCollectionIdentifier(), e);
            throw new RuntimeException("创建Milvus集合失败: " + e.getMessage(), e);
        }
        // 注意：这里不需要手动关闭客户端，因为 MilvusClientFactory 现在使用连接池管理
    }

    /**
     * 创建存储关联数据的集合
     * 
     * @param collection
     * @return
     */
    public void createMetaLinkCollection(DataDefinitionDto collection) {
        MilvusServiceClient milvusClient = null;
        try {
            milvusClient = milvusClientFactory.getClient(RagConstant.MILVUS_MARK);

            MilvusCollectionParamDto milvusCollectionParamDto = getDefaultCreateCollectionParam(
                    collection.getCollectionIdentifier() + RagConstant.META_POSTFIX, collection.getDescription(),
                    collection.getVectorModel());

            milvusClient.createCollection(milvusCollectionParamDto.getCreateCollectionParam());
            milvusClient.createIndex(milvusCollectionParamDto.getCreateIndexParam());
            
            log.info("成功创建Milvus元数据集合: {}", collection.getCollectionIdentifier() + RagConstant.META_POSTFIX);
        } catch (Exception e) {
            log.error("创建Milvus元数据集合失败: {}", collection.getCollectionIdentifier() + RagConstant.META_POSTFIX, e);
            throw new RuntimeException("创建Milvus元数据集合失败: " + e.getMessage(), e);
        }
        // 注意：这里不需要手动关闭客户端，因为 MilvusClientFactory 现在使用连接池管理
    }

    /**
     * 获取默认的创建集合参数
     * 
     * @param collectionName 集合名称
     * @param description    集合描述
     * @param vectorModel    向量模型标识
     * @return 创建集合参数
     */
    public MilvusCollectionParamDto getDefaultCreateCollectionParam(String collectionName,
            String description, String vectorModel) {
        // 创建字段类型
        FieldType idField = FieldType.newBuilder()
                .withName(RagConstant.LangchainMilvusFields.FIELD_ID)
                .withDataType(DataType.VarChar)
                .withMaxLength(64)
                .withPrimaryKey(true)
                .withAutoID(false)
                .build();

        FieldType embeddingField = FieldType.newBuilder()
                .withName(RagConstant.LangchainMilvusFields.FIELD_VECTOR)
                .withDataType(DataType.FloatVector)
                .withDimension(DimensionUtil.getDimension(vectorModel))
                .build();

        FieldType textField = FieldType.newBuilder()
                .withName(RagConstant.LangchainMilvusFields.FIELD_TEXT)
                .withDataType(DataType.VarChar)
                .withMaxLength(65535)
                .build();

        FieldType metadataField = FieldType.newBuilder()
                .withName(RagConstant.LangchainMilvusFields.FIELD_METADATA)
                .withDataType(DataType.JSON)
                .withMaxLength(65535)
                .build();

        /*
         * CollectionSchemaParam collectionSchemaParam =
         * CollectionSchemaParam.newBuilder()
         * .addFieldType(idField)
         * .addFieldType(embeddingField)
         * .addFieldType(textField)
         * .addFieldType(metadataField)
         * .build();
         */
        // 创建集合
        CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withDescription(description)
                .withShardsNum(2)
                .addFieldType(idField)
                .addFieldType(embeddingField)
                .addFieldType(textField)
                .addFieldType(metadataField)
                .build();

        // 创建索引
        CreateIndexParam createIndexParam = CreateIndexParam.newBuilder()
                .withCollectionName(collectionName)
                .withFieldName(RagConstant.LangchainMilvusFields.FIELD_VECTOR)
                .withIndexType(RagConstant.MILVUS_INDEX_TYPE)
                .withMetricType(RagConstant.MILVUS_METRIC_TYPE)
                .withExtraParam(RagConstant.MILVUS_EXTRA_PARAM)
                .build();

        MilvusCollectionParamDto milvusCollectionParamDto = MilvusCollectionParamDto.builder()
                .createCollectionParam(createCollectionParam)
                .createIndexParam(createIndexParam)
                .build();

        return milvusCollectionParamDto;
    }
}
